import logging
from typing import List, Dict, Optional
from concurrent.futures import ThreadPoolExecutor, as_completed
from .core_handler import CoreHandler
from dataclasses import dataclass
from threading import Lock
from .proxy_manager import ProxyManager
from .package import variables
import time

logger = logging.getLogger(__name__)

@dataclass
class UnlockTask:
    """解锁任务数据类"""
    record_id: int
    imei: str
    email: str
    first_name: str
    last_name: str
    phone: str
    is_test: bool = False

class MultiTaskHandler:
    def __init__(self, max_workers: int = 3):
        """
        初始化多任务处理器
        :param max_workers: 最大并行任务数
        """
        self.max_workers = max_workers
        self.tasks: List[UnlockTask] = []
        self.results: Dict[int, bool] = {}
        self.lock = Lock()
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        
    def add_task(self, task: UnlockTask) -> None:
        """
        添加解锁任务
        :param task: 解锁任务对象
        """
        with self.lock:
            self.tasks.append(task)
            
    def process_task(self, task: UnlockTask, task_id: int) -> bool:
        """
        处理单个解锁任务
        :param task: 解锁任务对象
        :param task_id: 任务ID
        :return: 是否成功
        """
        try:
            # 为每个任务创建新的代理管理器实例
            proxy_manager = ProxyManager()
            
            # 获取IP并确保与其他任务不同
            max_retries = 3
            retry_count = 0
            initial_ip = None
            
            while retry_count < max_retries:
                initial_ip = proxy_manager.refresh_ip()
                if not initial_ip:
                    logger.error(f"任务 {task_id}: 初始化代理失败，请检查代理配置")
                    return False
                
                # 检查是否有其他任务使用相同的IP
                ip_in_use = False
                for other_task_id in range(self.max_workers):
                    if other_task_id != task_id:
                        other_proxy_manager = getattr(variables, f'proxy_manager_{other_task_id}', None)
                        if other_proxy_manager and other_proxy_manager.current_ip == initial_ip:
                            ip_in_use = True
                            break
                
                if not ip_in_use:
                    break
                    
                logger.warning(f"任务 {task_id}: IP {initial_ip} 已被其他任务使用，重试获取新IP...")
                retry_count += 1
                time.sleep(2)  # 等待一段时间再重试
            
            if retry_count >= max_retries:
                logger.error(f"任务 {task_id}: 无法获取独立IP，已重试 {max_retries} 次")
                return False
            
            logger.info(f"任务 {task_id} 使用IP: {initial_ip}")
            
            # 创建处理器实例
            handler = CoreHandler(task_id=task_id)
            
            # 设置该任务的代理管理器
            setattr(variables, f'proxy_manager_{task_id}', proxy_manager)
            
            # 执行解锁流程
            success = handler.run(
                imei=task.imei,
                email=task.email,
                first_name=task.first_name,
                last_name=task.last_name,
                phone=task.phone
            )
            
            # 获取结果消息
            result_message = handler.state.result_message or handler.state.error_message or "未知错误"
            
            # 如果不是测试模式，发送结果到后台
            if not task.is_test:
                # 延迟导入以避免循环依赖
                from .main import send_result_to_backend
                send_result_to_backend(
                    record_id=task.record_id,
                    imei=task.imei,
                    success=success,
                    message=result_message
                )
            
            # 记录结果
            with self.lock:
                self.results[task.record_id] = success
                
            # 清理该任务的代理管理器
            if hasattr(variables, f'proxy_manager_{task_id}'):
                delattr(variables, f'proxy_manager_{task_id}')
                
            return success
            
        except Exception as e:
            error_msg = str(e)
            logger.error(f"任务 {task_id} 执行失败: {error_msg}")
            # 如果不是测试模式，发送错误到后台
            if not task.is_test:
                # 延迟导入以避免循环依赖
                from .main import send_result_to_backend
                # 获取CoreHandler中的错误消息（如果有）
                if hasattr(handler, 'state') and handler.state.error_message:
                    error_msg = handler.state.error_message
                send_result_to_backend(
                    record_id=task.record_id,
                    imei=task.imei,
                    success=False,
                    message=error_msg
                )
            with self.lock:
                self.results[task.record_id] = False
            
            # 清理该任务的代理管理器
            if hasattr(variables, f'proxy_manager_{task_id}'):
                delattr(variables, f'proxy_manager_{task_id}')
                
            return False
            
    def run_all_tasks(self) -> Dict[int, bool]:
        """
        执行所有任务
        :return: 任务结果字典，key为记录ID，value为是否成功
        """
        futures = []
        
        try:
            # 提交所有任务
            for i, task in enumerate(self.tasks):
                future = self.executor.submit(self.process_task, task, i)
                futures.append(future)
                
            # 等待所有任务完成
            for future in as_completed(futures):
                try:
                    future.result()
                except Exception as e:
                    logger.error(f"任务执行出错: {str(e)}")
                    
        finally:
            # 关闭线程池
            self.executor.shutdown(wait=True)
            
        return self.results
        
    def clear(self) -> None:
        """清理任务列表和结果"""
        with self.lock:
            self.tasks.clear()
            self.results.clear() 